import argparse

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='示例脚本：接收CUDA设备和端口参数')
    
    # 添加参数
    parser.add_argument(
        '-c',
        '--cuda_id', 
        type=str,
        required=True,  # 必须传入
        help='CUDA设备ID，例如 "0" 或 "0,1"（字符串类型）'
    )
    parser.add_argument(
        '-p',
        '--port', 
        type=int,
        required=True,  # 必须传入
        help='port 端口'
    )
    
    # 解析参数
    args = parser.parse_args()
    return args

args = parse_args()

import os
os.environ['CUDA_VISIBLE_DEVICES']=args.cuda_id

from flask import Flask, request, send_file
import json
import torch
import numpy as np
from io import BytesIO
from PIL import Image
from diffusers import FluxKontextPipeline
from utils.MODEL_CKP import FLUX_KONTEXT
from utils.util_flux import horizontal_concat_images, process_img_1024

# pipe = FluxKontextPipeline.from_pretrained(FLUX_KONTEXT, 
#                                            dtype=torch.bfloat16)
# pipe.to("cuda")
par_lora_dir = '/mnt/nas/shengjie/puthere_output_4_0922'
lora_dir = lambda id: f'{par_lora_dir}/checkpoint-{id}'
lora_name = 'pytorch_lora_weights.safetensors'
def get_lora_path(lora_dir , lora_name):
    return os.path.join( lora_dir , lora_name )
lora_path = get_lora_path( lora_dir(6000) , lora_name )
# pipe.load_lora_weights(lora_path)

# 固定 prompt
prompt = "[puthere] put it here"

from demo_sam import load_sam2,get_mask_by_img_coor_inputbox
from flask_cors import CORS  # 导入 CORS

app = Flask(__name__)
CORS(app)  # 允许所有域名访问（开发环境推荐）

predictor = load_sam2()

img, control_img, rec_coor = None,None,None
save_dir = '/mnt/nas/shengjie/datasets/puthere_data_process'

# 这是你的图像分割函数，这里只是一个示例
def segment_image(image, rect):
    # 示例：简单地在矩形区域绘制一个红色矩形
    # 实际应用中替换为你的分割算法

    # 确保 image 是 PIL.Image.Image 类型，如果不是则转换
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 将图像转换为OpenCV格式
    # img_array = np.array(image)
    # img_cv = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    
    # 绘制矩形
    x, y, w, h = int(rect['x']), int(rect['y']), int(rect['width']), int(rect['height'])
    # cv2.rectangle(img_cv, (x, y), (x + w, y + h), (0, 0, 255), 2)
    
    # 转换回PIL格式
    # result_img = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))

    mask = get_mask_by_img_coor_inputbox(predictor, image, x,y ,x+w,y+h)

    # res = Composite ( mask, image ) 

    return mask


def puthere_pipeline(pipeline,image,prompt=prompt):
    # 确保 image 是 PIL.Image.Image 类型，如果不是则转换
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    device = pipeline.device
    weight_dtype = pipeline.dtype
    autocast_ctx = torch.autocast(device.type, weight_dtype)

    with autocast_ctx:
        image = pipeline(
            prompt=prompt,
            image=image,
            num_inference_steps=8,
            guidance_scale=2.5,
            generator=torch.Generator(device=device).manual_seed(42),
            max_sequence_length=512,
            height=image.height,
            width=image.width,
        ).images[0]
    
    image = process_img_1024( '', image, image.size,)

    return image


@app.route('/crop', methods=['POST'])
def process_image():
    if 'image' not in request.files or 'rect' not in request.form:
        return "Missing image or rectangle data", 400
    
    # 获取上传的图像
    file = request.files['image']
    image = Image.open(file.stream)
    
    # 获取矩形坐标
    rect = json.loads(request.form['rect'])

    print('rect::>>', rect)
    
    # 调用分割函数
    result_image = segment_image(image, rect) # np.unique = 0 1 
    # INSERT_YOUR_CODE
    # 确保输出为PIL.Image类型
    if isinstance(result_image, np.ndarray):
        # INSERT_YOUR_CODE
        # 将0/1 mask转为RGBA格式的PIL mask
        # 0->全透明，1->全不透明白色
        mask = (result_image > 0).astype(np.uint8) * 255  # 0/1 -> 0/255
        # INSERT_YOUR_CODE
        # 把mask存入tmp_mask.png中查看效果
        # Image.fromarray(mask).save("tmp_mask.png")
        rgba = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
        # RGB三个通道用原图代替, 最终的图就是被遮罩的原图
        rgba[..., 0:3] = np.array(image.convert("RGB"))  # 用原图RGB通道
        rgba[..., 3] = mask   # A通道为mask
        result_image = Image.fromarray(rgba, mode='RGBA')
    
    # 将结果图像转换为字节流返回
    img_io = BytesIO()
    result_image.save(img_io, 'PNG')
    img_io.seek(0)

    global img, control_img, rec_coor
    img = image.convert('RGB')
    # control_img = mask==255的原图部分
    # 先获取mask（A通道），再用mask把原图抠出来
    if result_image.mode == 'RGBA':
        rgba_arr = np.array(result_image)
        mask = rgba_arr[..., 3]  # A通道
        rgb = np.array(image.convert("RGB"))
        control_img_arr = np.ones_like(rgb)*255
        control_img_arr[mask == 255] = rgb[mask == 255]
        control_img = Image.fromarray(control_img_arr, mode='RGB')
    else:
        control_img = result_image.convert('RGB')
    rec_coor = rect
    horizontal_concat_images([img, control_img]).save('tmp_put.jpg')
    
    return send_file(img_io, mimetype='image/png')

pipe = None
def load_model_once():
    global pipe
    if pipe is None:
        pipe = FluxKontextPipeline.from_pretrained(FLUX_KONTEXT, torch_dtype=torch.bfloat16)
        pipe.to("cuda")
        pipe.load_lora_weights(lora_path)
    return pipe

@app.route('/puthere', methods=['POST'])
def process_puthere():
    if 'image' not in request.files:
        return "Missing image or rectangle data", 400
    
    # 获取上传的图像
    file = request.files['image']
    image = Image.open(file.stream)

    global prompt
    pipe = load_model_once()
    result_img = puthere_pipeline(pipe, image, prompt)
        
    # 将结果图像转换为字节流返回
    img_io = BytesIO()
    result_img.save(img_io, 'PNG')
    img_io.seek(0)

    return send_file(img_io, mimetype='image/png')

if __name__ == '__main__':
    load_model_once()
    app.run(host='0.0.0.0',port=args.port, debug=True,use_reloader=False)

